package com.zijie.hosptal.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class RagConfig {

    private EmbeddingModel qwenEmbeddingModel;
    private EmbeddingStore<TextSegment> embeddingStore;

    @Autowired
    public void setEmbeddingStore(EmbeddingStore<TextSegment> embeddingStore) {
        this.embeddingStore = embeddingStore;
    }

    @Autowired
    public void setQwenEmbeddingModel(EmbeddingModel qwenEmbeddingModel) {
        this.qwenEmbeddingModel = qwenEmbeddingModel;
    }

    @Bean
    public ContentRetriever ragService() {
        //1.加载文档
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs");
        //2.文档切割 最大1000字符 每次最多重叠200字符
        DocumentByParagraphSplitter document =
                new DocumentByParagraphSplitter(800, 180);
        //3.自定义文档加载在器  文档转化为向量存储到数据库中
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(document)
                .textSegmentTransformer(textSegment -> TextSegment.from
                        (textSegment.metadata().getString("file_name") + "\n" + textSegment.text(), textSegment.metadata()))
                //使用向量模型
                .embeddingModel(qwenEmbeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        //4.自定义加载器
        ingestor.ingest(documents);
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)//存储
                .embeddingModel(qwenEmbeddingModel)//向量模型
                .maxResults(5)//最多返回5条数据
                .minScore(0.75)//最小分数 过滤条件
                .build();
        return contentRetriever;
    }
}
